# -*- coding: utf-8 -*-
"""
Created on 04/05/2023
@author: 曾导SJTU
"""

import numpy as np
from precision_data import fp
from kernels import *

# 温度修正
def conduction_coefs(case, dim, ncx, ncy, ncz, ncoef, dt, spht, con, heat_src, x, y, z, dens, T, uf, vf, wf, ct):

    # print('---------------')
    # print('conduction coefs')
    # print('---------------')

    # 定义与case对象相关的局部变量
    id_aP = case.id_aP
    id_aE = case.id_aE
    id_aW = case.id_aW
    id_aN = case.id_aN
    id_aS = case.id_aS
    if dim == 3:
        id_aT = case.id_aT
        id_aB = case.id_aB
    id_bsrc = case.id_bsrc
    idt = fp(1.0) / dt

    # 获取面积和体积
    dx = case.dx
    dy = case.dy
    dz = case.dz
    area_x, area_y, area_z, vol = cal_area_vol(dx, dy, dz)

    idx = fp(1.0) / dx
    idy = fp(1.0) / dy
    idz = fp(1.0) / dz
    
    #对每一个网格进行循环，计算系数aP, aE, aW, aN, aS, aT, aB, bsrc 
    for k in range(ncz):
        for j in range(ncy):
            for i in range(ncx):
                
                # 初始化系数，令系数为零
                aE = fp(0.0); aW = fp(0.0);
                aN = fp(0.0); aS = fp(0.0);
                aT = fp(0.0); aB = fp(0.0);
                aP = fp(0.0)
      
                # 初始化源和源系数为零（这里的热源没有引入，所以可以不考虑）
                sC    = heat_src    # 恒热源
                sP    = fp(0.0)
                bsrc  = fp(0.0)

                # 计算每个网格中的系数，aP, aE, aW, aN, aS, aT, aB, bsrc  
                # 按照Patankar的书计算系数
                # 东方向的面
                if(i==ncx-1):               # 如果是x方向处的最后一个网格
                    rho   = dens[i,j,k]     # dens是密度场
                else:
                    rho   = fp(0.5)*(dens[i,j,k]+dens[i+1,j,k])

                mul = con; mur = con        # con传导系数
                ul    = uf[i+1,j,k]         # uf：中间速度场
                ur    = uf[i+1,j,k]
                # 通过一个公式去计算这个系数，这个函数在kernels.py文件中定义
                aE    = a_nb(area_x, idx, ul, ur, mul, mur, rho, -fp(1.0))

                # 下面的内容同上
                # 西方向的面
                if(i==0):
                    rho   = dens[i,j,k]
                else:
                    rho   = fp(0.5)*(dens[i,j,k]+dens[i-1,j,k])

                mul = con; mur = con
                ul    = uf[i,j,k]
                ur    = uf[i,j,k]
                aW    = a_nb(area_x, idx, ul, ur, mul, mur, rho, fp(1.0))

                # 北方向的面
                if(j==ncy-1):
                    rho   = dens[i,j,k]
                else:
                    rho   = fp(0.5)*(dens[i,j,k]+dens[i,j+1,k])

                mul = con; mur = con
                ul  = vf[i,j+1,k]
                ur  = vf[i,j+1,k]
                aN  = a_nb(area_y, idy, ul, ur, mul, mur, rho, -fp(1.0))

                # 南方向的面
                if(j==0):
                    rho   = dens[i,j,k]
                else:
                    rho   = fp(0.5)*(dens[i,j,k]+dens[i,j-1,k])

                mul = con; mur = con
                ul  = vf[i,j,k]
                ur  = vf[i,j,k]
                aS  = a_nb(area_y, idy, ul, ur, mul, mur, rho, fp(1.0))

                if dim==3:

                    # 顶面
                    if(k==ncz-1):
                        rho   = dens[i,j,k]
                    else:
                        rho   = fp(0.5)*(dens[i,j,k]+dens[i,j,k+1])

                    mul = con; mur = con
                    ul  = wf[i,j,k+1]
                    ur  = wf[i,j,k+1]
                    aT  = a_nb(area_z, idz, ul, ur, mul, mur, rho, -fp(1.0))

                    # 底面
                    if(k==0):
                        rho   = dens[i,j,k]
                    else:
                        rho   = fp(0.5)*(dens[i,j,k]+dens[i,j,k-1])

                    mul = con; mur = con
                    ul  = wf[i,j,k]
                    ur  = wf[i,j,k]
                    aB  = a_nb(area_z, idz, ul, ur, mul, mur, rho, fp(1.0))

                # 对ap进行更新
                rho   = dens[i,j,k]
                aP0   = rho*spht*vol*idt
                aP    = aP + aE + aW + aN + aS + aT + aB + aP0 - sP*vol

                # 按照我们参考的教科书存储系数
                ct[i,j,k,id_aP]   = aP
                ct[i,j,k,id_aE]   = -aE
                ct[i,j,k,id_aW]   = -aW
                ct[i,j,k,id_aN]   = -aN
                ct[i,j,k,id_aS]   = -aS
                if dim==3:
                    ct[i,j,k,id_aT]  = -aT
                    ct[i,j,k,id_aB]  = -aB

                bsrc  = sC*vol + aP0*T[i,j,k]
                ct[i,j,k,id_bsrc] = bsrc
          
    return
    
# 对网格的边界类型做处理
def conduction_coef_bcs(case, fluidboundary, dim, ncx, ncy, ncz, ncoef, dt, con, x, y, z, dens, T, ct):

    # print('---------------')
    # print('conduction coef bcs')
    # print('---------------')

    # 定义与case对象相关的局部变量
    id_aP = case.id_aP
    id_aE = case.id_aE
    id_aW = case.id_aW
    id_aN = case.id_aN
    id_aS = case.id_aS
    if dim == 3:
        id_aT = case.id_aT
        id_aB = case.id_aB
    id_bsrc = case.id_bsrc
    
    xc = case.xc
    yc = case.yc
    if dim==3:
        zc = case.zc

    # 定义与fluidboundary对象相关的局部变量
    bcid  = fluidboundary.bcid

    fid_e = fluidboundary.fid_e
    fid_w = fluidboundary.fid_w
    fid_n = fluidboundary.fid_n
    fid_s = fluidboundary.fid_s
    if dim==3:
        fid_t = fluidboundary.fid_t
        fid_b = fluidboundary.fid_b
    
    bcs = fluidboundary.bcs
    bcs_temp = fluidboundary.bcs_temp
    
    bc_none   = fluidboundary.bc_none  
    bc_wall   = fluidboundary.bc_wall  
    bc_inlet  = fluidboundary.bc_inlet 
    bc_outlet = fluidboundary.bc_outlet
    
    temp_bc_constant = fluidboundary.temp_bc_constant
    temp_bc_heatflux = fluidboundary.temp_bc_heatflux

    # 温度BC系数
    # 对每一个网格进行循环
    for k in range(ncz):
        for j in range(ncy):
            for i in range(ncx):
      
                bcid_e = bcid[i,j,k,fid_e]
                bcid_w = bcid[i,j,k,fid_w]
                bcid_n = bcid[i,j,k,fid_n]
                bcid_s = bcid[i,j,k,fid_s]
                # 3D
                
                bc_e = bcs[bcid_e].type
                bc_w = bcs[bcid_w].type
                bc_n = bcs[bcid_n].type
                bc_s = bcs[bcid_s].type
                
                ttype_e = bcs_temp[bcid_e].temp_type
                ttype_w = bcs_temp[bcid_w].temp_type
                ttype_n = bcs_temp[bcid_n].temp_type
                ttype_s = bcs_temp[bcid_s].temp_type
                
                bc_te = bcs_temp[bcid_e].T
                bc_tw = bcs_temp[bcid_w].T
                bc_tn = bcs_temp[bcid_n].T
                bc_ts = bcs_temp[bcid_s].T
                
                bc_fluxe = bcs_temp[bcid_e].heat_flux
                bc_fluxw = bcs_temp[bcid_w].heat_flux
                bc_fluxn = bcs_temp[bcid_n].heat_flux
                bc_fluxs = bcs_temp[bcid_s].heat_flux

                # 东方向的面 BC
                if bc_e == bc_inlet:
                    ct[i,j,k,id_aP]   -= ct[i,j,k,id_aE]
                    ct[i,j,k,id_bsrc] -= 2.0 * ct[i,j,k,id_aE] * bc_te
                    ct[i,j,k,id_aE]   = 0.0
                elif bc_e == bc_outlet:
                    ct[i,j,k,id_aP]   += ct[i,j,k,id_aE]
                    ct[i,j,k,id_aE]   = 0.0
                elif bc_e == bc_wall:
                    if ttype_e == temp_bc_constant:
                        ct[i,j,k,id_aP]   -= ct[i,j,k,id_aE]
                        ct[i,j,k,id_bsrc] -= 2.0 * ct[i,j,k,id_aE] * bc_te
                        ct[i,j,k,id_aE]   = 0.0
                    elif ttype_e == temp_bc_heatflux:
                        ct[i,j,k,id_aP]   += ct[i,j,k,id_aE]
                        ct[i,j,k,id_bsrc] += ct[i,j,k,id_aE] * bc_fluxe * 2.0 * (x[i+1] - xc[i]) / con
                        ct[i,j,k,id_aE]   = 0.0

                # 西方向的面 BC
                if bc_w == bc_inlet:
                    ct[i, j, k, id_aP] -= ct[i, j, k, id_aW]
                    ct[i, j, k, id_bsrc] -= 2.0 * ct[i, j, k, id_aW] * bc_tw
                    ct[i, j, k, id_aW] = 0.0
                elif bc_w == bc_outlet:
                    ct[i, j, k, id_aP] += ct[i, j, k, id_aW]
                    ct[i, j, k, id_aW] = 0.0
                elif bc_w == bc_wall:
                    if ttype_w == temp_bc_constant:
                        ct[i, j, k, id_aP] -= ct[i, j, k, id_aW]
                        ct[i, j, k, id_bsrc] -= 2.0 * ct[i, j, k, id_aW] * bc_tw
                        ct[i, j, k, id_aW] = 0.0
                    elif ttype_w == temp_bc_heatflux:
                        ct[i, j, k, id_aP] += ct[i, j, k, id_aW]
                        ct[i, j, k, id_bsrc] += ct[i, j, k, id_aW] * bc_fluxw * 2.0 * (x[i] - xc[i]) / con
                        ct[i, j, k, id_aW] = 0.0
                
                # 北方向的面 BC
                if bc_n == bc_inlet:
                    ct[i, j, k, id_aP] -= ct[i, j, k, id_aN]
                    ct[i, j, k, id_bsrc] -= 2.0 * ct[i, j, k, id_aN] * bc_tn
                    ct[i, j, k, id_aN] = 0.0
                elif bc_n == bc_outlet:
                    ct[i, j, k, id_aP] += ct[i, j, k, id_aN]
                    ct[i, j, k, id_aN] = 0.0
                elif bc_n == bc_wall:
                    if ttype_n == temp_bc_constant:
                        ct[i, j, k, id_aP] -= ct[i, j, k, id_aN]
                        ct[i, j, k, id_bsrc] -= 2.0 * ct[i, j, k, id_aN] * bc_tn
                        ct[i, j, k, id_aN] = 0.0
                    elif ttype_n == temp_bc_heatflux:
                        ct[i, j, k, id_aP] += ct[i, j, k, id_aN]
                        ct[i, j, k, id_bsrc] += ct[i, j, k, id_aN] * bc_fluxn * 2.0 * (y[j+1] - yc[j]) / con
                        ct[i, j, k, id_aN] = 0.0

                # 南方向的面 BC
                if bc_s == bc_inlet:
                    ct[i,j,k,id_aP]   -= ct[i,j,k,id_aS]
                    ct[i,j,k,id_bsrc] -= 2.0*ct[i,j,k,id_aS]*bc_ts
                    ct[i,j,k,id_aS]    = 0.0
                elif bc_s == bc_outlet:
                    ct[i,j,k,id_aP]   += ct[i,j,k,id_aS]
                    ct[i,j,k,id_aS]    = 0.0
                elif bc_s == bc_wall:
                    if ttype_s == temp_bc_constant:
                        ct[i,j,k,id_aP]   -= ct[i,j,k,id_aS]
                        ct[i,j,k,id_bsrc] -= 2.0*ct[i,j,k,id_aS]*bc_ts
                        ct[i,j,k,id_aS]    = 0.0
                    elif ttype_s == temp_bc_heatflux:
                        ct[i,j,k,id_aP]   += ct[i,j,k,id_aS]
                        ct[i,j,k,id_bsrc] += ct[i,j,k,id_aS]*bc_fluxs*2.0*(y[j]-yc[j])/con
                        ct[i,j,k,id_aS]    = 0.0

                # 3D

    return

